#include "merge_sort.hpp"

#include <algorithm>

#include "mpi.h"

void merge_sort_recursive(std::vector<int64_t> &input,
                          std::vector<int64_t> &ret,
                          size_t start, size_t end) {
  if (start >= end) {
    return;
  }
  size_t mid = (start + end) / 2;
  merge_sort_recursive(input, ret, start, mid);
  merge_sort_recursive(input, ret, mid + 1, end);
  size_t i = start, j = mid + 1, k = start;
  while (i <= mid && j <= end) {
    if (input[i] < input[j]) {
      ret[k++] = input[i++];
    } else {
      ret[k++] = input[j++];
    }
  }
  while (i <= mid) {
    ret[k++] = input[i++];
  }
  while (j <= end) {
    ret[k++] = input[j++];
  }
  for (size_t i = start; i <= end; ++i) {
    input[i] = ret[i];
  }
}

std::vector<int64_t> MergeSortTester::Baseline(
    const std::vector<int64_t> &input) {
  std::vector<int64_t> input_copy(input);
  std::vector<int64_t> ret(input.size());
  merge_sort_recursive(input_copy, ret, 0, input.size() - 1);
  return ret;
}

std::vector<int64_t> MergeSortTester::Implementation(
    const std::vector<int64_t> &input) {
  MPI_Comm comm = MPI_COMM_WORLD;
  int rank, num_procs;
  MPI_Comm_rank(comm, &rank);
  MPI_Comm_size(comm, &num_procs);

  std::vector<int64_t> ret(input.size(), 0);
  // TODO: use MPI to speedup it!
  ret = input;
  std::sort(ret.begin(), ret.end());
  MPI_Finalize();
  return ret;
}
